import networkx as nx
import matplotlib.pyplot as plt
import random
import json
from networkx.readwrite import json_graph
import matplotlib

# 修正中文显示问题
matplotlib.rcParams['font.sans-serif'] = ['Microsoft YaHei']  # 中文字体
matplotlib.rcParams['axes.unicode_minus'] = False  # 处理负号问题


def generate_waxman_topology(num_nodes, alpha=0.5, beta=0.2):
    """
    使用 Waxman 模型生成具有空间约束的随机网络拓扑
    :param num_nodes: 节点数量
    :param alpha: 影响连接概率的参数（控制节点间的距离对连接概率的影响）
    :param beta: 控制边的吸引范围
    :return: 生成的图对象
    """
    # 生成 Waxman 模型图
    G = nx.waxman_graph(num_nodes, alpha, beta)

    # 添加节点属性
    for node in G.nodes():
        G.nodes[node]['type'] = 'center' if len(list(G.neighbors(node))) > 1 else 'leaf'

    # 添加边属性
    for u, v in G.edges():
        G.edges[u, v]['weight'] = random.randint(1, 10)  # 边的权重（可以是带宽、延迟等）

    return G


def visualize_topology(G, title="Waxman 随机网络拓扑图"):
    """
    可视化网络拓扑
    :param G: 图对象
    :param title: 图的标题
    """
    pos = nx.spring_layout(G)  # 使用 spring 布局
    node_colors = ['red' if G.nodes[node]['type'] == 'center' else 'lightblue' for node in G.nodes()]
    nx.draw_networkx_nodes(G, pos, node_color=node_colors, node_size=500)
    nx.draw_networkx_edges(G, pos, edge_color='gray')
    nx.draw_networkx_labels(G, pos, font_size=10, font_color='black')
    plt.title(title)
    plt.show(block=True)  # 显示图形并保持窗口打开


def export_topology_data(G, file_name="waxman_graph_data.json"):
    """
    导出拓扑数据到 JSON 文件
    :param G: 图对象
    :param file_name: 输出文件名
    """
    data = json_graph.node_link_data(G)
    with open(file_name, 'w') as f:
        json.dump(data, f, indent=4)
    print(f"拓扑数据已导出到 {file_name}")


def main():
    # 参数设置
    num_nodes = 50  # 总节点数
    alpha = 0.5  # 参数 alpha（控制节点间距离对连接的影响）
    beta = 0.2  # 参数 beta（控制边的吸引范围）

    # 生成 Waxman 拓扑
    G = generate_waxman_topology(num_nodes, alpha, beta)

    # 可视化拓扑
    visualize_topology(G)

    # 导出数据
    export_topology_data(G, "waxman_graph_data.json")


if __name__ == "__main__":
    main()
